package api.lock;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.junit.Test;

import java.io.IOException;
import java.util.List;
import java.util.Random;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;


/**
 * 使用zookeeper原生api，结合zookeeper的临时顺序节点和watcher机制，简单实现独占锁
 * （注意：测试时觉得watcher这块好像有问题，但又不确定？）
 * <p>
 * 独占锁实现思路：
 * 1、客户端连接zookeeper，并在/lock下创建临时的且有序的子节点，第一个客户端对应的子节点为/lock/lock-0000000000，
 * 第二个为/lock/lock-0000000001，以此类推。
 * <p>
 * 2、客户端获取/lock下的子节点列表，判断自己创建的子节点是否为当前子节点列表中序号最小的子节点，如果是则认为获得锁，
 * 否则监听/lock的子节点变更消息，获得子节点变更通知后重复此步骤直至获得锁；
 * <p>
 * 3、获取到锁之后，执行该线程的业务代码；
 * <p>
 * 4、删除锁对应的node节点，释放锁，watcher机制在删除节点时会触发，其他线程依次进行上述步骤；
 */
public class DistributeLock {

    //根节点
    private static final String ROOT_LOCKS = "/LOCKS";

    private ZooKeeper zooKeeper;

    //会话超时时间
    private int sessionTimeout;

    //记录锁节点id
    private String lockID;

    //节点的数据
    private final static byte[] data = {1, 2};

    private CountDownLatch countDownLatch = new CountDownLatch(1);

    public DistributeLock() throws IOException, InterruptedException {
        this.zooKeeper = ZookeeperClient.getInstance();
        this.sessionTimeout = ZookeeperClient.getSessionTimeout();
    }


    //获取锁的方法
    public boolean lock() {
        try {
            //LOCKS/00000001
            lockID = zooKeeper.create(ROOT_LOCKS + "/", data, ZooDefs.Ids.
                    OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);

            System.out.println(Thread.currentThread().getName() + "->成功创建了lock节点[" + lockID + "], 开始去竞争锁");

            List<String> childrenNodes = zooKeeper.getChildren(ROOT_LOCKS, true);//获取根节点下的所有子节点
            //排序，从小到大
            SortedSet<String> sortedSet = new TreeSet<String>();
            for (String children : childrenNodes) {
                sortedSet.add(ROOT_LOCKS + "/" + children);
            }

            String first = sortedSet.first(); //拿到最小的节点
            if (lockID.equals(first)) {
                //表示当前就是最小的节点
                System.out.println(Thread.currentThread().getName() + "->成功获得锁，lock节点为:[" + lockID + "]");
                return true;
            }

            SortedSet<String> lessThanLockId = sortedSet.headSet(lockID);
            if (!lessThanLockId.isEmpty()) {
                String prevLockID = lessThanLockId.last();//拿到比当前LOCKID这个几点更小的上一个节点
                zooKeeper.exists(prevLockID, new LockWatcher(countDownLatch));
                countDownLatch.await(sessionTimeout, TimeUnit.MILLISECONDS);
                //上面这段代码意味着如果会话超时或者节点被删除（释放）了
                System.out.println(Thread.currentThread().getName() + " 成功获取锁：[" + lockID + "]");
                return true;
            }
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        return false;
    }


    public boolean unlock() {
        System.out.println(Thread.currentThread().getName() + "->开始释放锁：[" + lockID + "]");
        try {
            zooKeeper.delete(lockID, -1);
            System.out.println("节点[" + lockID + "]成功被删除");
            return true;
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }

        return false;
    }


    /**
     * 单元测试测试多线程，注意使用CountDownLatch，让执行单元测试线程等待多个线程完成线程任务
     */
    @Test
    public void testDistributeLock() {
        final CountDownLatch latch = new CountDownLatch(10);
        Random random = new Random();
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                DistributeLock lock = null;
                try {
                    lock = new DistributeLock();
                    lock.lock();
                    Thread.sleep(random.nextInt(500));
                    latch.countDown();
                } catch (IOException | InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    if (lock != null) {
                        lock.unlock();
                    }
                }
            }).start();
        }

        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) throws InterruptedException {
        final CountDownLatch latch = new CountDownLatch(2);
        Random random = new Random();
        for (int i = 0; i < 2; i++) {
            new Thread(() -> {
                DistributeLock lock = null;
                try {
                    lock = new DistributeLock();
                    lock.lock();
                    Thread.sleep(10000);
                } catch (IOException | InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    if (lock != null) {
                        lock.unlock();
                    }

                    latch.countDown();
                }
            }).start();
        }

        latch.await();
    }


    @Test
    public void test() throws KeeperException, InterruptedException {
        String lockId = zooKeeper.create("/ceshi" + "/", data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        System.out.println(lockId);


        SortedSet<String> sortedSet = new TreeSet<>();
        sortedSet.add("2");
        sortedSet.add("1");
        sortedSet.add("3");
        sortedSet.add("4");

        System.out.println(sortedSet);

        //获取sortedSet集合中比3小的元素，组成集合
        SortedSet<String> lessThan3Set = sortedSet.headSet("3");
        System.out.println(lessThan3Set);

        DistributeLock lock = null;
        try {
            lock = new DistributeLock();
            lock.lock();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (lock != null) {
                lock.unlock();
            }
        }
    }
}
